'''
检索神经网络模型的中存在图像的结果，并且将其作为图像的预处理模型
'''
from models.segnet import segnet
from matplotlib import pyplot as plt
import matplotlib
import os 
import numpy as np
import sys
'''torch'''
import torch
import torch.nn as torch_nn
import torch.nn.functional as F
import torchvision.transforms as transforms
'''ourfunction'''
from models.get_models import get_models
from utils.utils import load_ckpt,getedge,getedge_sig
from datasetloader.get_datasetLoader import get_augmentations
from datasetloader.augmentations import BandVI,OneHot,Normalize,ToTensor
from datasetloader.datasetPath import getRootdirFromDatasetName
import math

'''
检测神经网络模型中的第一层conv矩阵的滤波效果作用
# 使用神经网络模型的条件约束
    1. 图像预预处理参数一致
    2. 神经网络模型的预处理加载

本文件代码，拟合功能：
    1.加载模型
    2.图像预处理方法集合

'''

# -------sobel滤波-----------

def sobelfilter(img_gray):
    '''定义sobel滤波初始算子'''
    kernel_x=np.array([[-1,0,1],[-2,0,2],[-1,0,1]])
    kernel_y=np.array([[1,2,1],[0,0,0],[-1,-2,-1]])
    kernel_x=torch.from_numpy(kernel_x).reshape(1,1,3,3).float()
    kernel_y=torch.from_numpy(kernel_y).reshape(1,1,3,3).float()
    # 考虑数据运算
    h,w=img_gray.shape
    img_gray=torch.from_numpy(img_gray).reshape(1,1,h,w).float()
    img_gray_ex=torch.nn.ReplicationPad2d(1)(img_gray)
    G_x=F.conv2d(img_gray_ex,kernel_x,padding=0)
    G_y=F.conv2d(img_gray_ex,kernel_y,padding=0)
    G=torch.pow(G_x,2)+torch.pow(G_y,2)
    G=torch.pow(G,0.5) # 求解梯度
    print(torch.max(G),torch.min(G))
    # 为避免div 0 这里进行数据规划
    G_mask=G>0 # 选择所有大于0 的值
    G=(G==0)+G # 所有大于0 的值+1 
    
    G_theta=G_y/G # 求解sin值，且最小值为0 
    G_theta=G_theta*G_mask # 将所有梯度值为0的值，置为0

    print(torch.max(G_theta),torch.min(G_theta))
    G_theta=torch.asin(G_theta) # 求解梯度方向 ,弧度制表达
    
    G=G.numpy().reshape(h,w)
    G_theta=G_theta.numpy().reshape(h,w)
    return [G,G_theta] # 返回梯度值与梯度方向


def NMS(G,G_theta):
    '''
    非极大值抑制，找寻像素点局部最大值，将非极大值点所对应的灰度值置0
    G： 梯度值
    G_theta: 梯度方向（弧度制）
    算法基本流程：
    Step 1 选中的x,y
    Step 2 选中的周围采样点 g1,g2,gn
    step 3 求解相交点 ptemp1 ptemp2
    step 4 求解相交点 pt1,pt2 的双线性差值
    step 5 比较 g_xy pt1 pt2 的值，如果g_xy 是最大值，则保留，不是则置为0
    step 6 重复过程 ，直至所有的可能性检索完成
    注：
    这里使用循环，因此可能整体性能会下降的巨大，因此这里改用基于torch的矩阵计算型，避免循环的使用
    '''
    h,w=G.shape
    G=torch.from_numpy(G).float() # 梯度值 转换成 tensor
    G_theta=torch.from_numpy(G_theta).float() # 梯度方向值 转换为tensor
    # 填充边界
    G_Ex=torch.nn.ReplicationPad2d(1)(G.reshape(1,1,h,w)).reshape(h+2,w+2) # 扩充梯度值边界
    G_theta_Ex=torch.nn.ReplicationPad2d(1)(G_theta.reshape(1,1,h,w)).reshape(h+2,w+2) # 扩充梯度方向值边界
    # 构建偏移矩阵
    offset_x=np.ones((h,w))*np.array(list(range(1,int(w+1)))) # 构建偏移矩阵 x 
    offset_y=np.ones((h,w))*np.array(list(range(1,int(h+1)))).transpose() # 构建偏移矩阵 y 
    offset_x=torch.from_numpy(offset_x).float() # 转换为 tensor
    offset_y=torch.from_numpy(offset_y).float() # 转换为 tensor
    # 求解梯度方向值，并构建梯度方向差值坐标 pt1,pt2。 其中pt2与pt1中心对称，假定中心为（0，0）,求解对应的插值情况
    # 分成四种情况（-pi/2 ,-pi/4],(-pi/4,pi/4],(pi/4,pi/2)，{-pi/2,pi/2} ,
    # 其中pi/2与-pi/2 在几何意义上一致，只讨论pi/2 的情况
    pt1_x=G_theta.clone()*0
    pt1_y=G_theta.clone()*0
    tan_theta=torch.tan(G_theta) # 计算所有角度的tan值，方便作为后续进行计算
    tan_theta_0=(tan_theta==0)*1+tan_theta
    tan_theta_reciprocal=1/tan_theta_0 # 获得所有零的倒数
    # 计算（-pi/2,-pi/4） 此时y=-1 ，x=-1/tan(theta)
    mask_theta=1.0*(G_theta<(math.pi/-4))*(G_theta>(math.pi/-2)) # 条件约束版
    temp_x= -1*mask_theta*tan_theta_reciprocal # 这里表示第四象限
    pt1_x=pt1_x+temp_x # -1/tan(theta)
    pt1_y=pt1_y+mask_theta*-1 # -1 
    # 计算 （-pi/4,pi/4） 此时x=1,y=tan(theta) 
    mask_theta=1.0*(torch.abs(G_theta)<math.pi/4) # mask 约束
    temp_y=mask_theta*tan_theta
    pt1_x=pt1_x+mask_theta # 1
    pt1_y=pt1_y+temp_y # tan(theta)
    # 计算 （pi/4,pi/2） 此时y=1,x=-1/tan(theta)
    mask_theta=1.0*(G_theta>math.pi/4)
    temp_x=mask_theta*tan_theta_reciprocal
    pt1_x=pt1_x+temp_x # 1/tan(theta)
    pt1_y=pt1_y+mask_theta # 1
    # 单独计算G_theta=pi/2 or -pi/2 的情况
    mask_theta=(torch.abs(G_theta)==math.pi/2)*1.0 # 获得所有的特殊值
    mask_theta_no_pi_div_2=(torch.abs(G_theta)<math.pi/2)*1.0 # 其余值
    pt1_x=pt1_x*mask_theta_no_pi_div_2 # 置0
    pt1_y=pt1_y*mask_theta_no_pi_div_2 # 置0
    pt1_x=pt1_x # 此时 x=0
    pt1_y=pt1_y+mask_theta # 此时 y=1
    # 求解对称阵 pt2
    pt2_x=pt1_x*-1 # x的对称
    pt2_y=pt1_y*-1 # y的对称
    # 偏差修正 [dx,dy]+[x,y]=[x+dx,y+dy] 偏差方程
    pt1_x=pt1_x+offset_x 
    pt1_y=pt1_y+offset_y
    pt2_x=pt2_x+offset_x
    pt2_y=pt2_y+offset_y
    # 偏差点 pt1
    # 偏差点 pt2
    # 求解采样点 pt1 的插值结果
    # 针对 pt1 这里采取使用的双线性插值
    # 双线性插值公式 
    # f(p)= b*d*Q11+a*d*Q21+b*c*Q12+a*c*Q22
    # a=x-x1
    # b=x2-x
    # c=y-y1
    # d=y2-y
    # Q11=(x1,y1)
    # Q12=(x1,y2)
    # Q21=(x2,y1)
    # Q22=(x2,y2)
    pt_x1=torch.floor(pt1_x) # 向下取整 获得 x1
    pt_x2=torch.ceil(pt1_x)  # 向上取整 获得 x2
    pt_y1=torch.floor(pt1_y) # 向下取整 获得 y1
    pt_y2=torch.ceil(pt1_y)  # 向上取整 获得 y2 
    # 构建计算系数 a，b,c,d
    a=pt1_x-pt_x1 # (x-x1)
    b=pt_x2-pt1_x # (x2-x)
    c=pt1_y-pt_y1 # (y-y1)
    d=pt_y2-pt1_y # (y2-y)
    a=a.reshape(-1) 
    b=b.reshape(-1)
    c=c.reshape(-1)
    d=d.reshape(-1)
    # 构建采样坐标
    pt_x1=pt_x1.numpy().astype(np.int32).reshape(-1) # x1 
    pt_x2=pt_x2.numpy().astype(np.int32).reshape(-1) # x2
    pt_y1=pt_y1.numpy().astype(np.int32).reshape(-1) # y1
    pt_y2=pt_y2.numpy().astype(np.int32).reshape(-1) # y2
    # 整体采样
    Q11=G_Ex[pt_x1,pt_y1].reshape(-1)
    Q12=G_Ex[pt_x1,pt_y2].reshape(-1)
    Q21=G_Ex[pt_x2,pt_y1].reshape(-1)
    Q22=G_Ex[pt_x2,pt_y2].reshape(-1)
    # 开始计算
    Fp1=b*d*Q11+a*d*Q21+b*c*Q12+a*c*Q22
    Fp1=Fp1.reshape(h,w)
    # 求解采样点 pt2 的插值结果 双线性插值--------------------------
    pt_x1=torch.floor(pt2_x) # 向下取整 获得 x1
    pt_x2=torch.ceil(pt2_x)  # 向上取整 获得 x2
    pt_y1=torch.floor(pt2_y) # 向下取整 获得 y1
    pt_y2=torch.ceil(pt2_y)  # 向上取整 获得 y2 
    # 构建计算系数 a，b,c,d
    a=pt2_x-pt_x1 # (x-x1)
    b=pt_x2-pt2_x # (x2-x)
    c=pt2_y-pt_y1 # (y-y1)
    d=pt_y2-pt2_y# (y2-y)
    a=a.reshape(-1) 
    b=b.reshape(-1)
    c=c.reshape(-1)
    d=d.reshape(-1)
    # 构建采样坐标
    pt_x1=pt_x1.numpy().astype(np.int32).reshape(-1) # x1 
    pt_x2=pt_x2.numpy().astype(np.int32).reshape(-1) # x2
    pt_y1=pt_y1.numpy().astype(np.int32).reshape(-1) # y1
    pt_y2=pt_y2.numpy().astype(np.int32).reshape(-1) # y2
    # 整体采样
    Q11=G_Ex[pt_x1,pt_y1].reshape(-1)
    Q12=G_Ex[pt_x1,pt_y2].reshape(-1)
    Q21=G_Ex[pt_x2,pt_y1].reshape(-1)
    Q22=G_Ex[pt_x2,pt_y2].reshape(-1)
    # 开始计算
    Fp2=b*d*Q11+a*d*Q21+b*c*Q12+a*c*Q22
    Fp2=Fp2.reshape(h,w)    
    # Fp1,Fp2 插值点，用于非极大值比较
    # 当p>Fp1 && p> Fp2 保留，其余置零
    Fp1=(Fp1<G)*1.0
    Fp2=(Fp2<G)*1.0
    Fp=Fp1+Fp2 
    Fp=(Fp>1)*1.0 # 取极大值，并置零其他值 
    G_NMS=G*Fp # 保留极大值
    G_NMS=G_NMS.numpy()
    return G_NMS

# ------------------







def LoadModel(modelName,modelcpk_path,num_cls,input_channels):
    '''
    加载模型
    args:
        modelName 需要参与计算的模型的名称
        modelcpk_path 训练好模型的保存路径
    return:
          model : 加载好对应权重的模型
    '''
    model=get_models(modelName,num_class=num_cls,channels_num=input_channels)
    epoch,model,best_iou=load_ckpt(model,None,modelcpk_path,'cuda')
    return epoch,model,best_iou

def preProcess(image,mean_arr,std_arr):
    '''
    图像 zeros 归一化
    '''
    image=(image-mean_arr)/std_arr # z-score 零-均值归一化
    return image 

# 执行部分，卷积核，网络，处理
def main_process_sig(modelinfo,imginfo):
    # modelinfo params
    modelName=modelinfo['modelname']
    modelcpk_path=modelinfo['modelcpk']
    num_cls=modelinfo['num_cls']
    input_channels=modelinfo['input_channels']
    # imginfo
    img=np.load(imginfo['img_path'])
    label=np.load(imginfo['label_path'])
    seg=np.load(imginfo['seg_path'])
    mean_arr=np.array(imginfo['mean'])
    std_arr=np.array(imginfo['std'])
    # 加载模型 
    epoch,model,best_iou=LoadModel(modelName,modelcpk_path,num_cls,input_channels)
    # 抽取第一层
    conv1=model.inc.conv.conv._modules['0'] # unet第一层结构

    img_process=preProcess(img,mean_arr,std_arr) # 图像预处理
    img_process=np.transpose(img_process,(2,0,1))

    c,h,w=img_process.shape   
    img_process=torch.from_numpy(img_process).reshape(1,c,h,w).float().cuda()  #  ([1, 3, 514, 514])
    # 抽取对应的核,经过卷积之后三通道影像，取出一个
    
    #img_process=torch.nn.ReplicationPad2d(1)(img_process)# 边界填充
    #kernel_weight=conv1.weight.detach().cpu().numpy()  # 所有核的权重值，3*3*3*64个
    # 输出到csv
    conv_img=model.inc(img_process)
    #conv_img=F.conv2d(img_process,conv1.weight,padding=0)    
    b,c,h,w=conv_img.shape   ## = ([1, 64, 512, 512])
    
    for i in range(c):
        # 输出波段运算值，找第一通道计算后的图
        
        plt.figure(figsize=(12,12)) # 图像指定宽和高
        plt.imshow(conv_img[0,i,:,:].detach().cpu().numpy())
        plt.axis('off')
        plt.tight_layout()
        plt.savefig('/media/gis/databackup/ayc/modellist/kernel_result/RGB234/band2/{}_{}.png'.format('unet',i),dpi=300)
        plt.pause(0.02)
        plt.close()
    print("over")


if __name__=='__main__':
    modelinfo={
        'modelname':'unet',
        'modelcpk':'/media/gis/databackup/ayc/model_hastrained_log/RGB234_unet_result/checkpoint/Bestest_mIou_ckpt_unet_dataName_RGB234.pth',
        'num_cls':2,
        'input_channels':3
    }
    imginfo={
        'img_path':'/media/gis/databackup/ayc/modellist/dataset/nanchang/E512/RGB234/test/img/2037.npy',
        'label_path':'/media/gis/databackup/ayc/modellist/dataset/nanchang/E512/RGB234/test/label/2037.npy',
        'seg_path':'/media/gis/databackup/ayc/modellist/dataset/nanchang/E512/RGB234/test/seg/2037.npy',
        'mean':[103.47446405573739, 89.24539776311728, 118.10080357709782],
        'std':[40.96571881718136, 36.232309241791654, 46.13046923546727]
    }
    main_process_sig(modelinfo,imginfo)
